package site.wanjiahao.map;

import java.util.Comparator;
import java.util.LinkedList;
import java.util.Queue;

public class TreeMap<K, V> implements Map<K, V> {

    private static final int BLACK = 1;

    private static final int RED = 0;

    // 容量大小
    protected int size;

    // 指向根节点
    protected Node root;

    private final Comparator<K> comparator;

    public TreeMap(Comparator<K> comparator) {
        this.comparator = comparator;
    }

    public TreeMap() {
        this(null);
    }


    public int size() {
        return size;
    }

    public boolean isEmpty() {
        return size == 0;
    }

    public void clear() {
        root = null;
        size = 0;
    }

    @Override
    public V put(K key, V value) {
        if (key == null) {
            throw new RuntimeException("key 不能为空");
        }

        if (root == null) {
            // 添加root节点
            root = new Node(key, value, null);
            size++;
            addAfter(root);
            return null;
        }

        Node cur = root;
        Node parent;
        int compareRes;
        if (comparator != null) {
            do {
                compareRes = comparator.compare(cur.key, key);
                if (compareRes > 0) {
                    parent = cur;
                    cur = cur.left;
                } else if (compareRes < 0) {
                    parent = cur;
                    cur = cur.right;
                } else {
                    // 覆盖处理，ele可能为引用类型，更加符合实际情况
                    V oldValue = cur.value;
                    cur.key = key;
                    cur.value = value;
                    return oldValue;
                }
            } while (cur != null);
        } else {
            do {
                compareRes = ((Comparable<K>) cur.key).compareTo(key);
                if (compareRes > 0) {
                    parent = cur;
                    cur = cur.left;
                } else if (compareRes < 0) {
                    parent = cur;
                    cur = cur.right;
                } else {
                    // 覆盖处理，ele可能为引用类型，更加符合实际情况
                    V oldValue = cur.value;
                    cur.key = key;
                    cur.value = value;
                    return oldValue;
                }
            } while (cur != null);
        }

        Node setNode = new Node(key, value, parent);
        if (compareRes > 0) {
            parent.left = setNode;
        } else {
            parent.right = setNode;
        }
        size++;
        addAfter(setNode);
        return null;
    }

    @Override
    public V get(K key) {
        Node node = node(key);
        return node != null ? node.value : null;
    }

    @Override
    public V remove(K key) {
        // 判断当前节点的度
        Node node = node(key);
        if (node == null) return null;

        if (hasTwoChildren(node)) {
            // 度为2：先找到对应的前驱（后继）节点，覆盖删除的节点，然后在删除对应的前驱（后继）节点（只存在度为1或者度为2）
            Node successor = successor(node);
            // predecessor不存在为空
            node.key = successor.key;
            node.value = successor.value;
            node = successor;
        }

        // 度为1：当前节点为父节点的什么节点，当前节点的子节点，作为这个节点存在
        Node replacement = node.left != null ? node.left : node.right;
        if (replacement != null) {
            replacement.parent = node.parent;
            // 度为1
        }
        if (node.isLeftNode()) {
            node.parent.left = replacement;
        } else if (node.isRightNode()) {
            node.parent.right = replacement;
        } else {
            root = replacement;
        }
        size--;
        removeAfter(node, replacement);
        return node.value;
    }

    @Override
    public boolean containsKey(K key) {
        return node(key) != null;
    }

    @Override
    public boolean containsValue(V value) {
        // 遍历完每一层，如何遍历完每一层，将读出的数据放入队列中
        Queue<Node> queue = new LinkedList<>();
        queue.offer(root);
        while (!queue.isEmpty()) {
            Node node = queue.poll();
            if (node.value == value) return true;
            if (node.left != null) {
                queue.offer(node.left);
            }

            if (node.right != null) {
                queue.offer(node.right);
            }
        }
        return false;
    }

    @Override
    public void traversal(Consumer<K, V> consumer) {
        inorderTraversal(consumer);
    }

    private void addAfter(Node node) {
        // 如果是根节点
        Node parent = node.parent;
        if (parent == null) {
            // 染色
            black(node);
            return;
        }

        // 如果父节点为黑色
        if (isBlack(parent)) {
            // 直接添加，不做处理
            return;
        }

        // 父节点是红色
        Node uncle = parent.sibling();
        // 祖父染红
        Node grand = red(parent.parent);
        if (isBlack(uncle)) {
            // uncle 不是红色 旋转处理
            if (parent.isLeftNode()) {
                if (node.isRightNode()) {
                    // LR
                    parent.rotateLeft();
                    parent = node;
                }
                // LL
                grand.rotateRight();
            } else {
                if (node.isLeftNode()) {
                    // RL
                    parent.rotateRight();
                    parent = node;
                }
                // RR
                grand.rotateLeft();
            }
            // 父亲染黑
            black(parent);
        } else {
            // uncle 是红色 上溢处理
            black(parent);
            black(uncle);
            addAfter(grand);
        }
    }

    private void removeAfter(Node node, Node replacement) {
        Node parent = node.parent;

        if (isRed(node)) {
            // 删除红色节点，不用做任何处理
            return;
        }

        if (replacement != null) {
            // 删除黑色节点，并且黑色节点存在替代节点
            // 替代节点染黑，其余不用做任何处理
            black(replacement);
            return;
        }

        // 删除的是根节点
        if (parent == null) return;

        // 删除单个黑色节点，做下移处理
        boolean left = parent.left == null || node.isLeftNode();
        Node sibling = left ? parent.right : parent.left;
        if (left) {
            // 当前节点在左边
            if (isRed(sibling)) {
                // 兄弟节点是红色
                black(sibling);
                red(parent);
                // 变成兄弟节点为黑节点情况
                parent.rotateLeft();
                sibling = parent.right;
            }
            if (sibling == null) return;
            // 兄弟节点是黑节点
            if (isRed(sibling.left) || isRed(sibling.right)) {
                if (isBlack(sibling.right)) {
                    // RL
                    sibling.rotateRight();
                    sibling = parent.right;
                }
                // RR
                parent.rotateLeft();
                // 染色
                color(sibling, colorOf(parent));
                color(sibling.left, RED);
                color(sibling.right, RED);
            } else {
                // 左右无节点，合并操作
                if (isBlack(parent)) {
                    red(sibling);
                    removeAfter(parent, null);
                } else {
                    black(parent);
                    red(sibling);
                }
            }
        } else {
            // 当前节点在右边
            if (isRed(sibling)) {
                // 兄弟节点是红色
                black(sibling);
                red(parent);
                // 变成兄弟节点为黑节点情况
                parent.rotateRight();
                sibling = parent.left;
            }
            if (sibling == null) return;
            // 兄弟节点是黑节点
            if (isRed(sibling.left) || isRed(sibling.right)) {
                if (isBlack(sibling.left)) {
                    // LR
                    sibling.rotateLeft();
                    sibling = parent.left;
                }
                // LL
                parent.rotateRight();
                // 染色
                color(sibling, colorOf(parent));
                color(sibling.left, RED);
                color(sibling.right, RED);
            } else {
                // 左右无节点，合并操作
                if (isBlack(parent)) {
                    red(sibling);
                    removeAfter(parent, null);
                } else {
                    black(parent);
                    red(sibling);
                }
            }
        }
    }

    // 前序遍历
    public void preorderTraversal(Consumer<K, V> consumer) {
        if (consumer == null) return;
        preorderTraversal(root, consumer);
    }

    public void preorderTraversal(Node node, Consumer<K, V> consumer) {
        // 1.递归出口
        if (node == null) return;
        if (consumer.isStop) return;
        // 外部自自定义消费逻辑
        consumer.isStop = consumer.accept(node.key, node.value);

        // 明确函数作用是什么，遍历当前拿到的节点
        preorderTraversal(node.left, consumer);
        preorderTraversal(node.right, consumer);
    }

    // 中序遍历
    public void inorderTraversal(Consumer<K, V> consumer) {
        if (consumer == null) return;
        inorderTraversal(root, consumer);
    }

    public void inorderTraversal(Node node, Consumer<K, V> consumer) {
        // 1.递归出口
        if (node == null || consumer.isStop) return;
        inorderTraversal(node.left, consumer);
        if (consumer.isStop) return;
        consumer.isStop = consumer.accept(node.key, node.value);
        // 明确函数作用是什么，遍历当前拿到的节点
        inorderTraversal(node.right, consumer);
    }

    // 后序遍历
    public void postorderTraversal(Consumer<K, V> consumer) {
        postorderTraversal(root, consumer);
    }

    public void postorderTraversal(Node node, Consumer<K, V> consumer) {
        // 1.递归出口
        if (node == null || consumer.isStop) return;

        // 明确函数作用是什么，遍历当前拿到的节点
        postorderTraversal(node.left, consumer);
        postorderTraversal(node.right, consumer);
        if (consumer.isStop) return;
        consumer.isStop = consumer.accept(node.key, node.value);
    }

    // 根据元素值，获取节点对象
    public Node node(K key) {
        Node node = root;
        if (comparator != null) {
            while (node != null) {
                int comp = comparator.compare(node.key, key);
                if (comp > 0) {
                    // 左边
                    node = node.left;
                } else if (comp < 0) {
                    // 右边
                    node = node.right;
                } else {
                    return node;
                }
            }
        } else {
            while (node != null) {
                int comp = ((Comparable<K>) node.key).compareTo(key);
                if (comp > 0) {
                    // 左边
                    node = node.left;
                } else if (comp < 0) {
                    // 右边
                    node = node.right;
                } else {
                    return node;
                }
            }
        }
        return null;
    }

    // 前驱节点
    private Node predecessor(Node node) {
        // 当前节点是否有左节点
        if (node.left != null) {
            node = node.left;
            while (node.right != null) {
                node = node.right;
            }
            // 循环终止，当前节点为前驱节点
            return node;
        } else if (node.parent != null) {
            // 父节点的右节点为当前节点，则为前驱节点
            while (node.parent != null && node.parent.right != node) {
                node = node.parent;
            }
            return node.parent;
        } else {
            // 左节点也为空，父节点为空，无前驱节点
            return null;
        }
    }

    // 后继节点
    private Node successor(Node node) {
        // 当前节点是否有左节点
        if (node.right != null) {
            node = node.right;
            while (node.left != null) {
                node = node.left;
            }
            // 循环终止，当前节点为前驱节点
            return node;
        } else if (node.parent != null) {
            // 父节点的右节点为当前节点，则为前驱节点
            while (node.parent != null && node.parent.left != node) {
                node = node.parent;
            }
            return node.parent;
        } else {
            // 左节点也为空，父节点为空，无前驱节点
            return null;
        }
    }

    private boolean hasTwoChildren(Node node) {
        return node.left != null && node.right != null;
    }

    private boolean nonChildren(Node node) {
        return node.left == null || node.right == null;
    }

    // 自定义工具方法
    private int colorOf(Node node) {
        return node == null ? BLACK : node.color;
    }

    private Node color(Node node, int color) {
        if (node != null) {
            node.color = color;
        }
        return node;
    }

    private boolean isBlack(Node node) {
        return colorOf(node) == BLACK;
    }

    private boolean isRed(Node node) {
        return colorOf(node) == RED;
    }

    private Node black(Node node) {
        return color(node, BLACK);
    }

    private Node red(Node node) {
        return color(node, RED);
    }

    private class Node {

        K key;

        V value;

        Node left;

        Node right;

        Node parent;

        // 默认添加红色节点，比较符合红黑树性质
        int color = RED;

        public Node(K key, V value, Node parent) {
            this.key = key;
            this.value = value;
            this.parent = parent;
        }

        public boolean isLeftNode() {
            return this.parent != null && (this.parent.left == this);
        }

        public boolean isRightNode() {
            return this.parent != null && (this.parent.right == this);
        }

        public Node sibling() {
            if (this != null) {
                if (isLeftNode()) {
                    return parent.right;
                }
                if (isRightNode()) {
                    return parent.left;
                }
            }
            return null;
        }

        public void rotateLeft() {
            Node next = this.right;
            this.right = next.left;
            next.left = this;

            rotateAfter(this, next, this.right);
        }

        public void rotateRight() {
            Node next = this.left;
            this.left = next.right;
            next.right = this;

            rotateAfter(this, next, this.left);
        }


        public void rotateAfter(Node cur, Node next, Node node) {
            Node parent = cur.parent;
            next.parent = parent;
            if (parent == null) {
                root = next;
            } else if (cur.isLeftNode()) {
                cur.parent.left = next;
            } else if (cur.isRightNode()) {
                cur.parent.right = next;
            }
            cur.parent = next;
            if (node != null) {
                node.parent = cur;
            }
        }
    }
}
